"""
本代码是头部姿态检测与识别的代码。将本代码直接保持在项目路径中，在自己的代码中引用该代码的函数即可使用。
代码中设置有抬头、摇头、低头（点头）的阈值角度，请根据实际需求修改这些阈值角度，以获得更好地检测与识别效果。（第83行至第106行代码）
"""
# -*- coding: utf-8 -*-
from imutils import face_utils
import numpy as np
import imutils
import dlib
import cv2
import math

def head_det(cap):
    # 世界坐标系(UVW)：填写3D参考点，该模型参考http://aifi.isr.uc.pt/Downloads/OpenGL/glAnthropometric3DModel.cpp
    object_pts = np.float32([[6.825897, 6.760612, 4.402142],  # 33左眉左上角
                             [1.330353, 7.122144, 6.903745],  # 29左眉右角
                             [-1.330353, 7.122144, 6.903745],  # 34右眉左角
                             [-6.825897, 6.760612, 4.402142],  # 38右眉右上角
                             [5.311432, 5.485328, 3.987654],  # 13左眼左上角
                             [1.789930, 5.393625, 4.413414],  # 17左眼右上角
                             [-1.789930, 5.393625, 4.413414],  # 25右眼左上角
                             [-5.311432, 5.485328, 3.987654],  # 21右眼右上角
                             [2.005628, 1.409845, 6.165652],  # 55鼻子左上角
                             [-2.005628, 1.409845, 6.165652],  # 49鼻子右上角
                             [2.774015, -2.080775, 5.048531],  # 43嘴左上角
                             [-2.774015, -2.080775, 5.048531],  # 39嘴右上角
                             [0.000000, -3.116408, 6.097667],  # 45嘴中央下角
                             [0.000000, -7.415691, 4.070434]])  # 6下巴角

    # 相机坐标系(XYZ)：添加相机内参
    K = [6.5308391993466671e+002, 0.0, 3.1950000000000000e+002,
         0.0, 6.5308391993466671e+002, 2.3950000000000000e+002,
         0.0, 0.0, 1.0]  # 等价于矩阵[fx, 0, cx; 0, fy, cy; 0, 0, 1]
    # 图像中心坐标系(uv)：相机畸变参数[k1, k2, p1, p2, k3]
    D = [7.0834633684407095e-002, 6.9140193737175351e-002, 0.0, 0.0, -1.3073460323689292e+000]

    # 像素坐标系(xy)：填写凸轮的本征和畸变系数
    cam_matrix = np.array(K).reshape(3, 3).astype(np.float32)
    dist_coeffs = np.array(D).reshape(5, 1).astype(np.float32)

    # 重新投影3D点的世界坐标轴以验证结果姿势
    reprojectsrc = np.float32([[10.0, 10.0, 10.0],
                               [10.0, 10.0, -10.0],
                               [10.0, -10.0, -10.0],
                               [10.0, -10.0, 10.0],
                               [-10.0, 10.0, 10.0],
                               [-10.0, 10.0, -10.0],
                               [-10.0, -10.0, -10.0],
                               [-10.0, -10.0, 10.0]])
    # 绘制正方体12轴
    line_pairs = [[0, 1], [1, 2], [2, 3], [3, 0],
                  [4, 5], [5, 6], [6, 7], [7, 4],
                  [0, 4], [1, 5], [2, 6], [3, 7]]


    def get_head_pose(shape):  # 头部姿态估计
        # （像素坐标集合）填写2D参考点，注释遵循https://ibug.doc.ic.ac.uk/resources/300-W/
        # 17左眉左上角/21左眉右角/22右眉左上角/26右眉右上角/36左眼左上角/39左眼右上角/42右眼左上角/
        # 45右眼右上角/31鼻子左上角/35鼻子右上角/48左上角/54嘴右上角/57嘴中央下角/8下巴角
        image_pts = np.float32([shape[17], shape[21], shape[22], shape[26],
                                shape[36],shape[39], shape[42], shape[45], shape[31],
                                shape[35],shape[48], shape[54], shape[57], shape[8]])
        # solvePnP计算姿势——求解旋转和平移矩阵：
        # rotation_vec表示旋转矩阵，translation_vec表示平移矩阵，cam_matrix与K矩阵对应，dist_coeffs与D矩阵对应。
        _, rotation_vec, translation_vec = cv2.solvePnP(object_pts, image_pts, cam_matrix, dist_coeffs)
        # projectPoints重新投影误差：原2d点和重投影2d点的距离（输入3d点、相机内参、相机畸变、r、t，输出重投影2d点）
        reprojectdst, _ = cv2.projectPoints(reprojectsrc, rotation_vec, translation_vec, cam_matrix, dist_coeffs)
        reprojectdst = tuple(map(tuple, reprojectdst.reshape(8, 2)))  # 以8行2列显示

        # 计算欧拉角calc euler angle
        rotation_mat, _ = cv2.Rodrigues(rotation_vec)  # 罗德里格斯公式（将旋转矩阵转换为旋转向量）
        pose_mat = cv2.hconcat((rotation_mat, translation_vec))  # 水平拼接，vconcat垂直拼接
        # decomposeProjectionMatrix将投影矩阵分解为旋转矩阵和相机矩阵
        _, _, _, _, _, _, euler_angle = cv2.decomposeProjectionMatrix(pose_mat)
        pitch, yaw, roll = [math.radians(_) for _ in euler_angle]
        pitch = math.degrees(math.asin(math.sin(pitch)))
        roll = -math.degrees(math.asin(math.sin(roll)))
        yaw = math.degrees(math.asin(math.sin(yaw)))
        print('pitch:{}, yaw:{}, roll:{}'.format(pitch, yaw, roll))
        return reprojectdst, euler_angle  # 投影误差，欧拉角


    # 定义常数
    # 抬头，若超过抬头角度阈值7次，则标注为抬头
    RAISE_THRESH = -8 #抬头角度阈值
    RAISE_AR_CONSEC_FRAMES = 7 #抬头次数阈值

    # 初始化帧计数器和抬头总数
    rCOUNTER = 0
    rTOTAL = 0

    # 摇头，若超过摇头角度阈值7次，则标注为摇头
    SHAKE_THRESH = 4.5 #摇头角度阈值
    SHAKE_AR_CONSEC_FRAMES = 7 #摇头次数阈值

    # 初始化帧计数器和摇头总数
    sCOUNTER = 0
    sTOTAL = 0

    # 点头，若超过点头角度阈值7次，则标注为点头
    NOD_THRESH = 0.3 #点头角度阈值
    NOD_AR_CONSEC_FRAMES = 7

    # 初始化帧计数器和点头总数
    hCOUNTER = 0
    hTOTAL = 0


    # 使用dlib.get_frontal_face_detector() 获得脸部位置检测器
    detector = dlib.get_frontal_face_detector()
    # 使用dlib.shape_predictor获得脸部特征位置检测器
    predictor = dlib.shape_predictor(
        'shape_predictor_68_face_landmarks.dat')

    # 分别获取左右眼面部标志的索引
    (lStart, lEnd) = face_utils.FACIAL_LANDMARKS_IDXS["left_eye"]
    (rStart, rEnd) = face_utils.FACIAL_LANDMARKS_IDXS["right_eye"]

    # 打开本地摄像头
    # cap = cv2.VideoCapture(0)

    # 从视频流循环帧
    while True:
        # 进行循环，读取图片，并对图片做维度扩大，并进灰度化
        ret, frame = cap.read()
        frame = imutils.resize(frame, width=720)
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        # 使用detector(gray, 0) 进行脸部位置检测
        rects = detector(gray, 0)

        # 循环脸部位置信息，使用predictor(gray, rect)获得脸部特征位置的信息
        for rect in rects:
            shape = predictor(gray, rect)

            # 将脸部特征信息转换为数组array的格式
            shape = face_utils.shape_to_np(shape)

            # 获取头部姿态
            reprojectdst, euler_angle = get_head_pose(shape)
            Yaw = euler_angle[2, 0]    # 取yaw旋转角度
            Pitch = euler_angle[0, 0]  # 取pitch旋转角度

            if Pitch < RAISE_THRESH:  # 角度小于阈值-8.0
                rCOUNTER += 1
            else:
                # 如果连续7次都小于阈值，则表示抬头一次
                if rCOUNTER >= RAISE_AR_CONSEC_FRAMES:  # 阈值：7
                    rTOTAL += 1
                # 重置点头帧计数器
                rCOUNTER = 0
            
            if abs(Yaw) < SHAKE_THRESH:  # 角度小于阈值4.5
                sCOUNTER += 1
            else:
                # 如果连续7次都小于阈值，则表示瞌睡点头一次
                if sCOUNTER >= SHAKE_AR_CONSEC_FRAMES:  # 阈值：7
                    sTOTAL += 1
                # 重置点头帧计数器
                sCOUNTER = 0
            
            if Pitch > NOD_THRESH:  # 角度大于阈值0.3
                hCOUNTER += 1
            else:
                # 如果连续7次都大于阈值，则表示瞌睡点头一次
                if hCOUNTER >= NOD_AR_CONSEC_FRAMES:  # 阈值：7
                    hTOTAL += 1
                # 重置点头帧计数器
                hCOUNTER = 0
            
            # 绘制正方体12轴
            for start, end in line_pairs:
                cv2.line(frame, (int(reprojectdst[start][0]),int(reprojectdst[start][1])),
                         (int(reprojectdst[end][0]),int(reprojectdst[end][1])), (0, 0, 255))
            # 显示角度结果
            cv2.putText(frame, "X: " + "{:7.2f}".format(euler_angle[0, 0]), (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.75,
                        (0, 255, 0), thickness=2)  # GREEN
            cv2.putText(frame, "Y: " + "{:7.2f}".format(euler_angle[1, 0]), (150, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.75,
                        (255, 0, 0), thickness=2)  # BLUE
            cv2.putText(frame, "Z: " + "{:7.2f}".format(euler_angle[2, 0]), (300, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.75,
                        (0, 0, 255), thickness=2)  # RED

            # 显示低抬头、摇头与点头次数
            cv2.putText(frame, "Raise: {}".format(rTOTAL), (450, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 0), 2)
            cv2.putText(frame, "Shake: {}".format(sTOTAL), (450, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 0), 2)
            cv2.putText(frame, "Nod: {}".format(hTOTAL), (450, 90), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 0), 2)

            # 进行画图操作，68个特征点标识
            for (x, y) in shape:
                cv2.circle(frame, (x, y), 1, (0, 0, 255), -1)

        # 按q退出
        cv2.putText(frame, "Press 'q': Quit", (20, 500), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (84, 255, 159), 2)
        # 窗口显示
        show2 = cv2.imshow("Frame", frame)
        return show2
